#include <rclcpp/rclcpp.hpp>
#include <sensor_msgs/msg/compressed_image.hpp>
#include <cv_bridge/cv_bridge.h>
#include <opencv2/opencv.hpp>
#include "task/facenet.h"
#include "task/retinaface.h"

#include <dirent.h>
#include <sys/stat.h>
#include <iostream>
#include <vector>
#include <string>
#include <cstring>


std::string getFileName(const std::string& path) {
    size_t pos = path.find_last_of("/\\");
    return (pos == std::string::npos) ? path : path.substr(pos + 1);
}

std::string getFileNameWithoutExtension(const std::string& path) {
    std::string fileName = getFileName(path);  // 使用上面的 getFileName 函数
    size_t lastDotPos = fileName.find_last_of('.');
    return (lastDotPos == std::string::npos) ? fileName : fileName.substr(0, lastDotPos);
}


std::vector<std::string> getDatFiles(const std::string& path) {
    std::vector<std::string> datFiles;
    DIR* dir = opendir(path.c_str());
    if (dir == nullptr) {
        std::cerr << "Could not open directory: " << path << std::endl;
        return datFiles;
    }

    struct dirent* entry;
    while ((entry = readdir(dir)) != nullptr) {
        std::string filename = entry->d_name;
        if (filename.length() >= 4 && 
            filename.substr(filename.length() - 4) == ".dat") {
            datFiles.push_back(path + filename);
        }
    }
    closedir(dir);
    return datFiles;
}

float eu_distance(float* input)
{
	float sum = 0;
	for (int i = 0; i < 128; ++i)
	{
		sum = sum + input[i] * input[i];
	}
	sum = sqrt(sum);
	return sum;
}

float compare_eu_distance(float* input1, float* input2)
{
	float sum = 0;
	for (int i = 0; i < 128; ++i)
	{
		sum = sum + (input1[i] - input2[i]) * (input1[i] - input2[i]);
	}
	sum = sqrt(sum);
	return sum;
}

float cos_similarity(float* input1, float* input2)
{
	float sum = 0;
	for (int i = 0; i < 128; ++i)
	{
		sum = sum + input1[i] * input2[i];
	}
	float tmp1 = eu_distance(input1);
	float tmp2 = eu_distance(input2);
	return sum / (tmp1 * tmp2);
}

void read_dat(std::string datPath, float* output)
{
    std::ifstream infile(datPath);
    std::string tmp;

    int i = 0;
    while (getline(infile, tmp))
    {
        output[i] = atof(tmp.c_str());
        i++;
    }
    infile.close();
}

class USBCameraSubscriber : public rclcpp::Node
{
public:
    USBCameraSubscriber() : Node("usb_camera_subscriber")
    {
        subscription_ = this->create_subscription<sensor_msgs::msg::CompressedImage>(
            "camera/image_raw/compressed", 30,
            std::bind(&USBCameraSubscriber::image_callback, this, std::placeholders::_1));

        /*加载模型权重*/
        retinaface_ = std::make_unique<RetinafaceCustom>();
        retinaface_->LoadModel("/mnt/sdcard/yp/face_recognition_demo/data/retinaFace_mobilenet640_fp.rknn");
        facenet_ = std::make_unique<facenetCustom>();
        facenet_->LoadModel("/mnt/sdcard/yp/face_recognition_demo/data/mobilefacenet.rknn");

        /*获取数据库中的dat列表*/
        datLists = getDatFiles("/mnt/sdcard/yp/ros2/ros_service_ws/src/image_service/datasets/");
        for (const auto& file : datLists) {
            std::cout << file << std::endl;
        }
    }

private:
    void image_callback(const sensor_msgs::msg::CompressedImage::SharedPtr msg) const
    {
        try
        {
            // 将压缩图像转换为OpenCV格式
            cv::Mat image = cv::imdecode(cv::Mat(msg->data), 1);

            retinaface_result result;
            retinaface_->Run(image, &result);

            int font = cv::FONT_HERSHEY_SIMPLEX;
            double scale = 2.0;
            cv::Scalar color(255, 0, 255);  // 白色
            int thickness = 2;

            // printf("result count : %d", result.count);
            for (int i = 0; i < result.count; ++i) {
                int x1 = result.object[i].box.left;
                int y1 = result.object[i].box.top;
                int x2 = result.object[i].box.right;
                int y2 = result.object[i].box.bottom;
                cv::Rect bbox(x1, y1, x2-x1, y2-y1);

                cv::Mat crop_img = image(bbox);
                cv::Mat align_img;
                retinaface_->FaceAlignment(crop_img, align_img, &result.object[i]);

                float res_vector[128];
                facenet_->Run(align_img, res_vector);

                float tmp2, tmp3;

                for (const auto& file : datLists) {
                     /*加载目标对象的脸部矢量*/
                    float target_vector[128];
                    read_dat(file, target_vector);

                    std::string fileNameWithoutExt = getFileNameWithoutExtension(file);

                    tmp2 = compare_eu_distance(res_vector, target_vector);
                    tmp3 = cos_similarity(res_vector, target_vector);


                    if(tmp3 > 0.7)
                    {
                        RCLCPP_INFO(this->get_logger(), u8"%s, eu_distance: %f, cos_similarity: %f.", fileNameWithoutExt.c_str(), tmp2, tmp3);
                        cv::rectangle(image, cv::Point(x1, y1), cv::Point(x2, y2), cv::Scalar(255, 0, 255), 3);
                        // 在图像上添加文本
                        cv::putText(image, fileNameWithoutExt, cv::Point(x1, y1), font, scale, color, thickness);
                        break;
                    }    
                }
            }

            // 显示图像
            cv::imshow("Received Image", image);
            cv::waitKey(1);  // 允许OpenCV处理窗口事件
        }
        catch (cv_bridge::Exception& e)
        {
            RCLCPP_ERROR(this->get_logger(), "无法转换图像: %s", e.what());
        }
    }

    rclcpp::Subscription<sensor_msgs::msg::CompressedImage>::SharedPtr subscription_;
    std::unique_ptr<RetinafaceCustom> retinaface_;
    std::unique_ptr<facenetCustom> facenet_;
    std::vector<std::string> datLists;
};

int main(int argc, char * argv[])
{
    rclcpp::init(argc, argv);
    rclcpp::spin(std::make_shared<USBCameraSubscriber>());
    rclcpp::shutdown();
    return 0;
}
